import os
import argparse
from pathlib import Path

'''
python train.py --gpu 0 -d summe_canonical -s original \
    --suffix devvv --discriminator_lr 1.5e-6 --lr 1.2e-4 \
    --discriminator_slow_start 12 --clip 6 --debug --fold-list 1 3 5
'''
parser = argparse.ArgumentParser("train cmd")

parser.add_argument('--gpu', type=str, required=True, help="which card to use")
parser.add_argument('-d', '--dataset', type=str, required=True,
                    choices=['dev', 'ytb_canonical', 'ovp_canonical', 'tvsum_canonical', 'summe_canonical',
                             'ytb_transfer', 'ovp_transfer', 'tvsum_transfer', 'summe_transfer',
                             'tvsum_aug', 'summe_aug', 'ytb_aug', 'ovp_aug'], help="Dataset")
parser.add_argument('-s', '--shift', type=str, required=True,
                    choices=['original', 'rep', 'sup', 'dpp'],
                    help="Variants of sparsity loss")
parser.add_argument('--discriminator_lr', type=float, default=1e-5)
parser.add_argument('--lr', type=float, default=1e-4)
parser.add_argument('--clip', type=float, default=5.0)    # 梯度裁剪
parser.add_argument('--discriminator_slow_start', type=int, default=15)
parser.add_argument('--suffix', type=str, required=True, help="suffix")
parser.add_argument('--n_epochs', type=int, default=30)
# parser.add_argument('--use-cpu', action='store_true', help="use cpu device")
parser.add_argument('--debug', action='store_true')
parser.add_argument('--eval', action='store_true')
parser.add_argument('--fold-list', nargs='+', type=int, help='run fold list', default=[1, 2, 3, 4, 5])

args = parser.parse_args()
print(args)

project_dir = Path(__file__).resolve().parents[1]
trainpy = project_dir.joinpath('train.py')
evalpy = project_dir.joinpath('evaluate_code/my_eval.py')
# TODO lr的格式控制:.2e
args.suffix = f'dslr{args.discriminator_lr:.2e}_lr{args.lr:.2e}_clip{args.clip}_slow{args.discriminator_slow_start}_{args.suffix}'
logdir = project_dir.parents[0].joinpath(
    f'data/results/{args.dataset}_{args.shift}_{args.suffix}')
try:
    logdir.mkdir()  # If the path already exists, FileExistsError is raised.
except Exception as e:
    if args.debug is False:
        raise e
print(f'train.py: {trainpy}')
print(f'logdir: {logdir}')

for fold in args.fold_list:
    out = logdir.joinpath(f'f{fold}.log')
    cmd = (f'CUDA_VISIBLE_DEVICES={args.gpu} python {trainpy} '
           f'--dataset {args.dataset} '
           f'--shift {args.shift} '
           f'--fold {fold} '
           f'--discriminator_slow_start {args.discriminator_slow_start} '
           f'--discriminator_lr {args.discriminator_lr} '
           f'--lr {args.lr} '
           f'--clip {args.clip} '
           f'--suffix {args.suffix} '
           f'--n_epochs {args.n_epochs} ')
    if args.debug is False:
        cmd += f'> {out} 2>&1 '
    os.system(cmd)

if args.eval:
    os.system(f'python {evalpy} --root {logdir}')
